import torch
from typing import Dict


def preprocess_batch(batch: Dict[str, torch.Tensor], device: torch.device = None, dtype: torch.dtype = None) -> Dict[str, torch.Tensor]:
    batch["img"] = batch["img"].to(device = device, dtype = dtype)
    return batch


def postprocess_batch(batch: Dict[str, torch.Tensor], device: torch.device = None, dtype: torch.dtype = None) -> Dict[str, torch.Tensor]:
    img: torch.Tensor = batch["img"][-1]
    img[:, 0] = torch.sum(img[:, :10], dim = 1)
    img[:, 2] = torch.sum(img[:, 10:], dim = 1)
    img[:, 1] = 0
    batch["img"] = img[:, :3].to(device = device, dtype = dtype) * 16
    return batch